from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
import os
import sys
# 添加父级目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.insert(0, parent_dir)
from my_common import load_auto_model, load_auto_tokenizer, load_flag_model

# 集合 杰卡德相似性，A B的交集大小除以A B的并集大小
def jaccard_similarity(sentence1, sentence2):
    set1 = set(sentence1.split(" "))
    set2 = set(sentence2.split(" "))
    intersection = set1.intersection(set2)
    union = set1.union(set2)
    return len(intersection)/len(union)

# 欧式距离，传统几何学距离计算
def euclidean_distance(A, B):
    print(A, B)
    print("euclidean distance:")
    # 计算方法
    dist = torch.sqrt(torch.sum(torch.pow(torch.subtract(A, B), 2), dim=-1))
    print(dist.item())

    # 或者直接使用torch cdist 函数
    print(torch.cdist(A, B, p=2).item())

# 余弦相似性
def cosine_similarity(A, B):
    # 原生实现
    # Compute the dot product of A and B
    dot_prod = sum(a*b for a, b in zip(A[0], B[0]))

    # Compute the magnitude of A and B
    A_norm = torch.sqrt(sum(a*a for a in A[0]))
    B_norm = torch.sqrt(sum(b*b for b in B[0]))
    cos_1 = dot_prod / (A_norm * B_norm)
    print("cosine similarity:")
    print(cos_1.item())

    # pytroch 实现
    res = torch.mm(A / A.norm(dim=1), B.T / B.norm(dim=1))
    print(res.item())

    # pytroch 方法调用
    res = F.cosine_similarity(A, B)
    print(res.item())

    # 内积/点积计算
    dot_prod = A @ B.T
    print(dot_prod.item())

    # BGE系列模型已经将输出的embedding向量归一化为1，因此使用点积和余弦相似度也会得到相同的结果
    flag_model = load_flag_model(query_instruction_for_retrieval='Represent this sentence for searching relevant passages:')
    sentence = "I am very interested in natural language processing"
    embedding = torch.tensor(flag_model.encode(sentence))
    print(torch.norm(embedding).item())


if __name__ == '__main__':
    s1 = "Hawaii is a wonderful place for holiday"
    s2 = "Peter's favorite place to spend his holiday is Hawaii"
    s3 = "Anna enjoys baking during her holiday"

    print(jaccard_similarity(s1, s2))
    print(jaccard_similarity(s1, s3))

    # 生成两个张量
    A = torch.randint(1, 7, (1, 4), dtype=torch.float32)
    B = torch.randint(1, 7, (1, 4), dtype=torch.float32)
    # 欧式距离
    euclidean_distance(A,B)
    #  余弦相似性
    cosine_similarity(A, B)